from pathlib import Path

import cv2
import numpy as np
import torch
from tqdm import tqdm

from model import mobilenet_v1
from torchvision.transforms import transforms

from dataLoader import ToTensor, Normalize, LP300WDataSet
from utity import _load, save


def fix_pose(dfa_fp='/home/chenxianhao/3DDFA_linux/3DDFA-master/3DDFA-master/train.configs/',
             img_fp='/home/chenxianhao/train_aug_120x120/',
             dst='/home/chenxianhao/FaceRecon3dM/dataset/'):
    train_names = Path(dfa_fp + 'train_aug_120x120.list.train').read_text().strip().split('\n')
    val_names = Path(dfa_fp + 'train_aug_120x120.list.val').read_text().strip().split('\n')
    print(f'文件名-训练集：{len(train_names)}, 验证集：{len(val_names)}')

    meta = _load(dst + 'dfa/param_mean_std_62d_120x120.pkl')
    param_mean = meta.get('mean')
    param_std = meta.get('std')

    checkpoint = torch.load('./results/3ddfaV2.pth', map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, 'mobilenet_1')(num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)
    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        kc = k.replace('module.', '')
        if kc in model_dict.keys():
            model_dict[kc] = checkpoint[k]
        if kc in ['fc_param.bias', 'fc_param.weight']:
            model_dict[kc.replace('_param', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    model.eval()

    model = model.cuda()

    old_meta = _load(dst + 'mean_std.pkl')
    old_train = _load(dst + 'param_train_normal.pkl')
    old_val = _load(dst + 'param_val_normal.pkl')
    param_mean_o = old_meta.get('mean')
    param_std_o = old_meta.get('std')

    param_train = []
    i = 0
    transform = transforms.Compose([ToTensor(), Normalize(mean=127.5, std=128)])
    with tqdm(total=len(train_names), desc='Train_param: ') as pbar:
        for s in train_names:
            image = cv2.imread(img_fp + s)
            img = cv2.resize(image, dsize=(120, 120), interpolation=cv2.INTER_LINEAR)
            input = transform(img).unsqueeze(0)
            with torch.no_grad():
                input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(np.float32)
                param = param * param_std + param_mean
                old_p = old_train[i]
                old_p = old_p * param_std_o + param_mean_o
                new_p = old_p
                new_p[:12] = param[:12]
                param_train.append(new_p)
            i += 1
            pbar.update(1)

    param_val = []
    i = 0
    with tqdm(total=len(val_names), desc='Val_param: ') as pbar:
        for s in val_names:
            image = cv2.imread(img_fp + s)
            img = cv2.resize(image, dsize=(120, 120), interpolation=cv2.INTER_LINEAR)
            input = transform(img).unsqueeze(0)
            with torch.no_grad():
                input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(np.float32)
                param = param * param_std + param_mean
                old_p = old_val[i]
                old_p = old_p * param_std_o + param_mean_o
                new_p = old_p
                new_p[:12] = param[:12]
                param_val.append(new_p)
            i += 1
            pbar.update(1)

    save(dst + 'param_train_n', param_train)
    save(dst + 'param_val_n', param_val)


def normalize(fp='dataset/'):
    """归一化参数 mean normalization

    Args:
        fp: 目录
    Returns:
        param_train_normal.pkl, param_val_normal.pkl, mean_std
    """
    param_train = _load(fp + 'param_train_n.pkl')
    param_train = np.array(param_train)
    mean = param_train.mean(0)
    std = param_train.std(0)
    param_train -= mean
    param_train /= std

    param_val = _load(fp + 'param_val_n.pkl')
    param_val -= mean
    param_val /= std

    save(fp + 'param_train_normal_fixed', param_train)
    save(fp + 'param_val_normal_fixed', param_val)
    save(fp + 'mean_std_fixed', {'mean': mean, "std": std})


if __name__ == '__main__':
    # fix_pose()
    normalize()
    print("finish")
